import gradio as gr
import json
import os
from pathlib import Path

from gxl_ai_utils.utils import utils_file

big_info_dict = {}
info_path = "./datasets.json"
# 加载之前的信息
if os.path.exists(info_path):
    big_info_dict = utils_file.load_dict_from_json(info_path)

def validate_input(data):
    """输入验证"""
    if not data["dataset_name"].strip():
        raise ValueError("数据集名称不能为空")
    if not data["data_list_path_lab"].strip():
        raise ValueError("data.list路径（lab server）不能为空")
    if data['task'] in big_info_dict and data['dataset_name'] in big_info_dict[data['task']]:
        raise ValueError("数据集名称已存在")
    return data

def save_dataset(task, dataset_name, description, data_list_path_lab, shards_list_path_lab, shards_list_path_huawei, duration):
    """保存/更新数据集"""
    data = {
        "task": task,
        "dataset_name": dataset_name,
    }
    info_dict = {
        "description": description,
        "data_list_path_lab": data_list_path_lab,
        "shards_list_path_lab": shards_list_path_lab,
        "shards_list_path_huawei": shards_list_path_huawei,
        "duration": duration
    }

    try:
        if data['task'] not in big_info_dict:
            big_info_dict[data['task']] = {}
        if dataset_name in big_info_dict[data['task']]:
            # 更新数据集
            big_info_dict[data['task']][dataset_name].update(info_dict)
        else:
            # 新增数据集
            big_info_dict[data['task']][dataset_name] = info_dict
        # 保存数据集信息
        utils_file.write_dict_to_json(big_info_dict, info_path)
        return f"保存成功: {dataset_name}"
    except Exception as e:
        return f"错误: {str(e)}"

def delete_dataset(dataset_name):
    """删除数据集"""
    try:
        if len(big_info_dict) == 0:
            return "无数据集"
        for task_tag, dataset_dict in big_info_dict.items():
            if dataset_name in dataset_dict:
                dataset_dict.pop(dataset_name)
                # 保存数据集信息
                utils_file.write_dict_to_json(big_info_dict, info_path)
                return f"删除成功: {dataset_name}"
        return f"数据集不存在: {dataset_name}"
    except Exception as e:
        return f"删除失败: {str(e)}"

def get_all_tasks():
    """获取所有数据类型"""
    tasks = big_info_dict.keys()
    return sorted(tasks) if tasks else ["无数据类型"]

def get_datasets_by_task(task):
    """根据类型获取数据集"""
    if task == "无数据类型":
        return []

    datasets = big_info_dict.get(task, {}).keys()
    return sorted(datasets) if datasets else ["无数据集"]

def load_dataset_info(dataset_name):
    """加载数据集完整信息"""
    if dataset_name in ["无数据集", "无数据类型"]:
        return None

    for task_tag, dataset_dict in big_info_dict.items():
        if dataset_name in dataset_dict:
            info_dict = dataset_dict[dataset_name]
            info_dict["task"] = task_tag
            info_dict["dataset_name"] = dataset_name
            return info_dict
    return {}

# 界面布局
with gr.Blocks(title="OSUM数据集管理系统") as app:
    gr.Markdown("# OSUM数据集管理系统")

    with gr.Tab("录入"):
        with gr.Row():
            with gr.Column():
                task = gr.Textbox(label="数据类型 (task)")
                dataset_name = gr.Textbox(label="数据集名称", )
                description = gr.Textbox(label="数据集介绍", lines=3)
            with gr.Column():
                data_list_path_lab = gr.Textbox(label="data.list路径（lab server）*", )
                shards_list_path_lab = gr.Textbox(label="shards_list.txt路径（lab server）")
                shards_list_path_huawei = gr.Textbox(label="shards_list.txt路径（huawei server）")
                duration = gr.Number(label="数据集时长（tar包个数）")
        save_btn = gr.Button("保存")
        save_output = gr.Textbox()

    with gr.Tab("删除"):
        del_name = gr.Textbox(label="输入要删除的数据集名称")
        del_btn = gr.Button("删除")
        del_output = gr.Textbox()

    with gr.Tab("查询/修改"):
        with gr.Row():
            with gr.Column():
                # 添加加载按钮
                load_button = gr.Button("加载数据类型")
                task_dropdown = gr.Dropdown(
                    label="选择数据类型",
                    interactive=True
                )
                dataset_dropdown = gr.Dropdown(
                    label="选择数据集",
                    interactive=True
                )
            with gr.Column():
                info_task = gr.Textbox(label="数据类型",)
                info_name = gr.Textbox(label="数据集名称",)
                info_desc = gr.Textbox(label="数据集介绍", lines=3)  # 修改为多行输入框，宽度更宽
                info_data_path_lab = gr.Textbox(label="data.list路径（lab server）",)
                info_shards_path_lab = gr.Textbox(label="shards_list.txt路径（lab server）",)
                info_shards_path_huawei = gr.Textbox(label="shards_list.txt路径（huawei server）",)
                info_duration = gr.Number(label="数据集时长",)

        with gr.Row():
            update_button = gr.Button("更新")
            update_output = gr.Textbox()

        # 点击加载按钮时更新task下拉框
        def load_tasks():
            return gr.Dropdown(choices=get_all_tasks(), value=None)

        load_button.click(
            fn=load_tasks,
            outputs=task_dropdown
        )

        # 当点击数据类型下拉框时动态更新选项
        def update_task_dropdown():
            return gr.Dropdown(choices=get_all_tasks())

        task_dropdown.select(
            fn=update_task_dropdown,
            outputs=task_dropdown
        )

        # 当数据类型变化时更新数据集下拉框
        def update_dataset_dropdown(task):
            return gr.Dropdown(
                choices=get_datasets_by_task(task),
                value=None
            )

        task_dropdown.change(
            fn=update_dataset_dropdown,
            inputs=task_dropdown,
            outputs=dataset_dropdown
        )

        # 当数据集变化时更新显示信息
        def update_info_panel(dataset_name):
            if not dataset_name or dataset_name in ["无数据集", "无数据类型"]:
                return [None] * 7

            info = load_dataset_info(dataset_name)
            if not info:
                return [None] * 7

            return [
                info["task"], info["dataset_name"], info["description"],
                info["data_list_path_lab"], info["shards_list_path_lab"], info["shards_list_path_huawei"], info["duration"]
            ]

        dataset_dropdown.change(
            fn=update_info_panel,
            inputs=dataset_dropdown,
            outputs=[info_task, info_name, info_desc, info_data_path_lab, info_shards_path_lab, info_shards_path_huawei, info_duration]
        )

        # 更新按钮点击事件
        def update_dataset_info(task, dataset_name, description, data_list_path_lab, shards_list_path_lab, shards_list_path_huawei, duration):
            if not dataset_name or dataset_name in ["无数据集", "无数据类型"]:
                return "数据集不存在"

            info = load_dataset_info(dataset_name)
            if not info:
                info = {}
                info["task"] = task
                info["dataset_name"] = dataset_name
                info["description"] = description
                info["data_list_path_lab"] = data_list_path_lab
                info["shards_list_path_lab"] = shards_list_path_lab
                info["shards_list_path_huawei"] = shards_list_path_huawei
                info["duration"] = duration
                # 保存更新后的数据集信息
                if info["task"] not in big_info_dict:
                    big_info_dict[info["task"]] = {}
                big_info_dict[info["task"]][dataset_name] = info
                utils_file.write_dict_to_json(big_info_dict, info_path)
                return "数据集不存在,直接写入该数据项"

            # 更新数据集信息
            info["task"] = task
            info["description"] = description
            info["data_list_path_lab"] = data_list_path_lab
            info["shards_list_path_lab"] = shards_list_path_lab
            info["shards_list_path_huawei"] = shards_list_path_huawei
            info["duration"] = duration

            # 保存更新后的数据集信息
            big_info_dict[info["task"]][dataset_name] = info
            utils_file.write_dict_to_json(big_info_dict, info_path)

            return "修改成功"

        update_button.click(
            fn=update_dataset_info,
            inputs=[info_task, info_name, info_desc, info_data_path_lab, info_shards_path_lab, info_shards_path_huawei, info_duration],
            outputs=update_output
        )

    # 其他事件处理
    save_btn.click(
        save_dataset,
        [task, dataset_name, description, data_list_path_lab, shards_list_path_lab, shards_list_path_huawei, duration],
        save_output
    )
    del_btn.click(
        delete_dataset,
        del_name,
        del_output
    )

if __name__ == "__main__":
    app.launch(
        share=True,
        server_name="0.0.0.0"
    )